# encoding=utf-8
from nltk.corpus import movie_reviews
from nltk.corpus import stopwords         #停用词
from sklearn.feature_extraction.text import TfidfVectorizer

pos_ids = movie_reviews.fileids('pos')
neg_ids = movie_reviews.fileids('neg')

pos_reviews = []
neg_reviews = []
for pids in pos_ids:
        pos_reviews.append((movie_reviews.raw(pids), 'positive'))
for nids in neg_ids:
        neg_reviews.append((movie_reviews.raw(nids), 'negative'))
pos_reviews=list(map(lambda x:x[0],pos_reviews))
neg_reviews=list(map(lambda x:x[0],neg_reviews))
with open('./data/train.txt', 'a') as file:
    for i in pos_reviews + neg_reviews:
        file.write(i + '\n')
    file.flush()
    file.close()
print('file write finished')
stopWords_list = list(stopwords.words('english'))  #英文停用分词集合
# print("英文停用分词集合：", stopWords_list)
def gettfidf(reviews):
    corpus = reviews
    # print(corpus)
    vector = TfidfVectorizer(stop_words=stopWords_list)
    tf_idf = vector.fit_transform(corpus)
    # print(tf_idf)
    word_list = vector.get_feature_names()      # 获取词袋模型的所有词
    weight_list = list(tf_idf.toarray())
    word_weight=[]
    for i in range(len(word_list)):
        x=[]
        weight=0
        for j in range(len(weight_list)):
            weight+=weight_list[j][i]
        x.append(word_list[i])
        x.append(weight)
        word_weight.append(x)
    return word_weight
pos_word_weight=gettfidf(pos_reviews)
neg_word_weight=gettfidf(neg_reviews)

pos_word_weight_sorted=sorted(pos_word_weight,key=(lambda x:x[1]),reverse=True)
neg_word_weight_sorted=sorted(neg_word_weight,key=(lambda x:x[1]),reverse=True)

pos_word=set()
neg_word=set()
for i in range(500):
    pos_word.add(pos_word_weight_sorted[i][0])
for i in range(500):
    neg_word.add(neg_word_weight_sorted[i][0])

print('pos:',len(pos_word-neg_word),pos_word-neg_word)
print('neg:',len(neg_word-pos_word),neg_word-pos_word)
with open('./data/sentiment_words.txt', 'a') as file:
    for i in pos_word-neg_word:
        file.write(i + '\t' + 'pos' + '\n')
    for i in neg_word-pos_word:
        file.write(i + '\t' + 'neg' + '\n')
    file.flush()
    file.close()